package com.dl4jlpr;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;

import ai.onnxruntime.OrtSession;
import lombok.extern.slf4j.Slf4j;
import org.datavec.image.loader.NativeImageLoader;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;
import org.nd4j.linalg.dataset.api.preprocessor.NormalizerStandardize;
import org.nd4j.linalg.indexing.INDArrayIndex;
import org.nd4j.onnxruntime.runner.OnnxRuntimeRunner;
import org.opencv.core.Core;

import java.io.File;
import java.util.*;

@Slf4j
public class LprNetOnnx {
    //private static final OrtEnvironment env = OrtEnvironment.getEnvironment();
    private static char[] chars={'京', '沪', '津', '渝', '冀', '晋', '蒙', '辽', '吉', '黑',
            '苏', '浙', '皖', '闽', '赣', '鲁', '豫', '鄂', '湘', '粤',
            '桂', '琼', '川', '贵', '云', '藏', '陕', '甘', '青', '宁',
            '新',
            '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
            'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K',
            'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V',
            'W', 'X', 'Y', 'Z', 'I', 'O', '-'
            };
    public static void detection() throws Exception{
        String path = ClassLoader.getSystemResource("").getPath();
        File f = new File (path+"/model/onnx/Final_LPRNet_model_cpu.onnx");
        NativeImageLoader originalLoad = new NativeImageLoader(24, 94, 3);
        INDArray image = originalLoad.asMatrix(new File(path+"/test/浙BR7S63.jpg")).add(-127.5).mul(0.0078125);
        //根据LprNet python源码
        /*
        def transform(self, img):
        img = img.astype('float32')
        img -= 127.5
        img *= 0.0078125
        img = np.transpose(img, (2, 0, 1))
        return img 修改了JPG归一化算法
         */
        //DataNormalization scaler = new ImagePreProcessingScaler(-0.9961,0.9961);//////仿照python进行归一化,python 的归一化算法没有细看,该值是根据推理终点矩阵倒推
        //DataNormalization scaler =new NormalizerStandardize
        //scaler.transform(image);
        //log.info("image sharp:{},rank:{}",Arrays.toString(image.shape()),image.rank());

        OnnxRuntimeRunner onnxRuntimeRunner = OnnxRuntimeRunner.builder().modelUri(f.getAbsolutePath()).build();
        Map<String,INDArray> inputs = new LinkedHashMap<>();
        inputs.put("input",image);
        Map<String, INDArray> exec = onnxRuntimeRunner.exec(inputs);
        INDArray out = exec.get("out");
        //out.reshape(1,68,18);
        //log.info("Out sharp:{},rank:{}",Arrays.toString(out.shape()),out.rank());
        //log.info("Out array:{}",out.data().asFloat());
        //log.info("Out data:{}",Arrays.toString(out.data().asFloat()));
        float[] o=out.data().asFloat();
        //String r=charsToString(GetMaxRowsArray(18,68,o),67);
        //log.info("License plate:{}",r);
        String r=charsToString(GetMaxColsArray(18,68,o),67 );
        log.info("Deduction license plate:{}",r);
    }
    public static int[] GetMaxRowsArray(int times,int len,float[] source) throws Exception
    {
        float max=0;
        int[] r=new int[times];
        for (int i=0;i<source.length;i++) {
            if (i%len==0) { // 当前值大于最大值，赋值为最大值
                max = source[i];
            }
            if(source[i]>max) {
                max=source[i];
                r[i/len]=i%len;
            }
        }
        //log.info(Arrays.toString(r));
        return r;
    }
    public static int[] GetMaxColsArray(int times,int len,float[] source) throws Exception
    {
        float max=0;
        int[] r=new int[times];
        int j;
        for (int i=0;i<times;i++) {
            for(j=0;j<len;j++) {
                if (source[j * times+i] > max) {
                    max = source[j * times+i];
                    r[i] = j;
                }
            }
            max=0;
        }
        log.info(Arrays.toString(r));
        return r;
    }
    public static String charsToString(int[] index,int emptyindex) throws Exception  /////去重转码
    {
        StringBuffer r=new StringBuffer("");
        int j=0;
        int pre=index[0];
        if(pre!=emptyindex) r.append(chars[pre]);
        for(int i:index)
        {
            if(i==emptyindex ||i==pre){
                if (i== emptyindex) pre=emptyindex;
                continue;
            }
            r.append(chars[i]);
            pre=i;
        }
        return String.valueOf(r);
    }
    public static void main(String[] args) throws Exception {
        detection();
    }
}
